from array import array
from random import choice
#import psyco; psyco.full()

def chararray(str):
    return array('c', str)

dna = chararray("")
rna = chararray("")

dnap = 0

#the numbers corresponding to the ?s and !s in patterns
pqueue = []
#the tuples and numbers corresponding to the ||s and n sub l's in templates
tqueue = []

#dna is an array.array('c', ...)
def execute(_dna):
    global dna, pqueue
    dna = _dna
    dnap = 0
    try:
        i = 0
        #while 1:
        for i in range(10):
            print "iteration %s" % i
            print "dna = %s (%s bytes)" % (dna[dnap:dnap+10], len(dna)-dnap)
            pqueue = []
            p = pattern()
            print "pattern %s %s" % (p, pqueue)
            t = template()
            print "template %s %s" % (t, tqueue)
            matchreplace(p, t)
            print "len(rna) %s" % (len(rna))
    except StopIteration: pass

def finish():
    f = file("rna.out", "w")
    f.write("".join(rna))
    f.close()
    #print "finish() dna <%s> rna <%s>" % (dna, rna)
    raise StopIteration()

def pattern():
    global dna, rna, pqueue, dnap

    p = chararray("")
    lvl = 0
    while 1:
        if dnap >= len(dna):
            finish()
        elif dna[dnap] == "C":
            dnap += 1
            p.append("I")
        elif dna[dnap] == "F":
            dnap += 1
            p.append("C")
        elif dna[dnap] == "P":
            dnap += 1
            p.append("F")
        elif dna[dnap:dnap+2] == chararray("IC"):
            dnap += 2
            p.append("P")
        elif dna[dnap:dnap+2] == chararray("IP"):
            dnap += 2
            #skip nat() DNA entries
            p.append("!")
            pqueue.insert(0, nat())
        elif dna[dnap:dnap+2] == chararray("IF"):
            dnap += 3
            #search for a DNA seq
            p.append("?")
            pqueue.insert(0, consts())
        elif dna[dnap:dnap+3] == chararray("IIP"):
            dnap += 3
            lvl += 1
            p.append("(")
        elif dna[dnap:dnap+3] in (chararray("IIC"), chararray("IIF")):
            dnap += 3
            if lvl == 0:
                return p
            else:
                lvl -= 1
                p.append(")")
        elif dna[dnap:dnap+3] == chararray("III"):
            rna.append(dna[dnap+3])
            dnap += 10
        else:
            #I or II are all that remain on the stack
            finish()

def nat():
    global dna, dnap
    #imperative version of nat

    try:
        p = dna[dnap:].index("P")
    except ValueError:
        finish()
    head = dna[dnap:dnap+p]
    dnap = dnap + p +1
    sum = 0
    for c in reversed(head):
        sum *= 2
        if c == "C": sum += 1
    return sum

def consts():
    global dna, dnap

    i = iter(dna[dnap:])
    res = chararray("") 
    deletions = 0
    try:
        for c in i:
            if c == "C":
                deletions += 1
                res.append("I")
            elif c == "F":
                deletions += 1
                res.append("C")
            elif c == "P":
                deletions += 1
                res.append("F")
            elif c == "I":
                d = i.next()
                if d == "C":
                    deletions += 2
                    res.append("P")
                else:
                    break
    except StopIteration:
        return res
    dnap += deletions
    return res

def template():
    global dna, rna, tqueue, dnap

    t = chararray("")
    while 1:
        if dnap >= len(dna):
            finish()
        if dna[dnap] == "C":
            dnap += 1
            t.append("I")
        elif dna[dnap] == "F":
            dnap += 1
            t.append("C")
        elif dna[dnap] == "P":
            dnap += 1
            t.append("F")
        elif dna[dnap:dnap+2] == chararray("IC"):
            dnap += 2
            t.append("P")
        elif dna[dnap:dnap+2] in (chararray("IP"), chararray("IF")):
            dnap += 2
            t.append(",")
            tqueue.insert(0, (nat(), nat()))
        elif dna[dnap:dnap+3] in (chararray("IIC"), chararray("IIF")):
            dnap += 3
            return t
        elif dna[dnap:dnap+3] == chararray("IIP"):
            dnap += 3
            #I'm not sure if |n| is supposed to actually operate on n
            t.append("|")
            tqueue.insert(0, nat())
        elif dna[dnap:dnap+3] == chararray("III"):
            rna.append(dna[dnap+3])
            dnap += 10
        else:
            #I or II
            finish()

def matchreplace(pat, t):
    global dna, dnap

    i = dnap
    e = [] #the environment is an array of DNA arrays
    c = [] #the stack of "(" pops
    for p in pat:
        if p in ("I", "C", "F", "P"):
            if dna[i] == p: i+= 1
            else: 
                print "unsuccessful match"
                return
        if p == "!":
            i += pqueue.pop()
            if i > len(dna): 
                print "unsuccessful match"
                return
        if p == "?":
            s = pqueue.pop()
            #XXX: this is likely to be too slow; profile.
            try:
                i = i + dna[i+1:].tostring().index(s.tostring())
            except ValueError:
                return
            except TypeError:
                import pdb; pdb.set_trace()
        if p == "(":
            #I've reversed the manual's order here and at ")" for efficiency
            c.append(i)
        if p == ")":
            e.append(dna[c.pop():i])
    #XXX: what do they mean by length?
    print "successful match of length %s %s" % (len(dna) - i, i)
    for i_, e_ in enumerate(e):
        print "e[%s] = %s... (%s bases)" % (i_, e_[:10].tostring(), len(e_))
    dnap = i
    replace(t,e)

def replace(tpl, e):
    global dna, dnap
    r = chararray("")
    for t in tpl:
        if t in ("I", "C", "F", "P"):
            r.append(t)
        elif t == ",":
            l, n = tqueue.pop()
            r.extend(protect(l, e[n]))
        elif t == "|":
            n = tqueue.pop()
            r.extend(asnat(len(e[n])))
    r.extend(dna[dnap:])
    dna = r
    dnap = 0

def protect(l, d):
    for i in range(l):
        d = quote(d)
    return d

def quote(ds):
    #XXX: some sort of string.translate might be quicker?
    #XXX: this is gonna be really slow - profile to see if it's a bottleneck
    #XXX: hashtable as a translate()? Real answer is a custom C type. This is
    #       gonna get ugly. Pyrex? Psyco?
    res = chararray("")
    for d in ds:
        if d == "I": res.append("C")
        elif d == "C": res.append("F")
        elif d == "F": res.append("P")
        elif d == "P": res.extend("IC")
        else: break
    return res

def asnat(n):
    res = chararray("")
    while 1:
        if n == 0: 
            res.append("P")
            break
        if n % 2 == 0:
            res.append("I")
        else:
            res.append("C")
        #won't work in py3k :)
        n /= 2
    return res


#################################
# TESTS
#################################

def setdna(pat):
    global dna, tqueue, pqueue, dnap
    dna = chararray(pat)
    rna = []
    #clear the stacks
    tqueue = []
    pqueue = []
    dnap = 0

def testmatchreplace(_dna, _pqueue, expected):
    setdna(_dna)
    global pqueue, dna
    pat = pattern()
    tem = template()
    matchreplace(pat, tem)
    if dna != expected:
        print "matchreplace failed! Expected <%s>, got <%s>" % (expected, dna)
        raise AssertionError()

def testmatchreplaces():
    tests = [
        #I know this DNA string produces this pattern with this pqueue
        ("IIPIPICPIICICIIF", "(!)P", [2]),
    ]
    for t in tests:
        testmatchreplace(*t)

def testasnat(n, expected):
    got = asnat(n)
    if got != chararray(expected):
        print "asnat broke %s" % locals()
        raise AssertionError()

def testexecute(pat, expected):
    global dna
    setdna(pat)
    execute(chararray(pat))
    if dna[dnap:] != chararray(expected):
        print "execute broke %s" % locals()
        print "dna <%s> rna <%s> pqueue <%s> tqueue <%s>" % (dna[dnap:], rna, pqueue,
            tqueue)
        raise AssertionError()

def testexecutes():
    tests = [
        ("IIPIPIICPIICIICCIICFCFC", "I"),
        # this state must be hit along the way (how to test?)
        #("IIPIPICPIICICIIFICCIFCCCPPIICCFPC", "PIICCFCFFPC"),
        ("IIPIPICPIICICIIFICCIFCCCPPIICCFPC", ""),
        # this state must be hit along the way (how to test?)
        #("IIPIPICPIICICIIFICCIFPPIICCFPC", "PICFC"),
        ("IIPIPICPIICICIIFICCIFPPIICCFPC", ""),
    ]
    for t in tests:
        testexecute(*t)

def testasnats():
    tests = [
        (0, "P"),
        (1, "CP"),
        (2, "ICP"),
        (3, "CCP"),
        (4, "IICP"),
        (5, "CICP"),
        (6, "ICCP"),
        (7, "CCCP"),
        (8, "IIICP"),
        (9, "CIICP"),
    ]
    for t in tests:
        testasnat(*t)

def testprotect(l, d, expected):
    got = protect(l, chararray(d))
    if got != chararray(expected):
        print "Protect broke %s" % locals()
        raise AssertionError()

def testprotects():
    tests = [
        (0, "", ""),
        (99, "", ""),
        (0, "ICFP", "ICFP"),
        (1, "ICFP", "CFPIC"),
        (2, "ICFP", "FPICCF"),
    ]
    for t in tests:
        testprotect(*t)

def testquote(d, expected):
    if quote(d) != chararray(expected):
        print "Quote broke d <%s> expected <%s>" % (d, expected)
        raise AssertionError()

def testquotes():
    tests = {
        "": "",
        "I": "C",
        "C": "F",
        "F": "P",
        "P": "IC",
        "ICFP": "CFPIC"
    }
    for d, expected in tests.iteritems():
        testquote(d, expected)

def testtemplate(pat, expected):
    global tqueue

    setdna(pat)
    if not isinstance(expected[0], str):
        try:
            template()
            print "template did not call finish() on pat <%s>" % pat
            raise AssertionError()
        except StopIteration: pass
    else:
        try:
            t = template()
            if t != chararray(expected[0]) or tqueue != expected[1] :
                print "Error in template()! expected <%s> got <%s> on dna: <%s>" \
                    % (expected, (t, tqueue), pat)
                raise AssertionError()
        except:
            print "unexpected exception! pat <%s> expected <%s> dna <%s>" \
                % (pat, expected, dna)
            raise AssertionError()

def testtemplates():
    tests = {
        "": (StopIteration, []),
        "C": (StopIteration, []),
        "F": (StopIteration, []),
        "IF": (StopIteration, []),
        "IIC": ("", []),
        "CIIC": ("I", []),
        "FIIC": ("C", []),
        "PIIC": ("F", []),
        "ICIIC": ("P", []),
        "IFICCPICPIIF": (",", [(6,2)]), #IF means "get two nats()", ICCP=6, ICP=2,
                                      #IIF signals end of seq
        "IIPICCPIIF": ("|", [6]),
        "IIIFFFFFFFIIC": ("", []),
        "IFICCPICPIIPICCPCIIF": (",|I", [6,(6,2)]),
    }
    for pat, expected in tests.iteritems():
        testtemplate(pat, expected)


def testconsts(pat, expected):
    setdna(pat)
    c = consts()
    if c != chararray(expected):
        print "Error in consts! expected <%s> got <%s> on dna: <%s>" \
            % (expected, c, pat)
        raise AssertionError()

def testconstss():
    tests = {
        "": "",
        "C": "I",
        "F": "C",
        "P": "F",
        "IC": "P",
        "IF": "",
        "IP": "",
        "II": "",
        "ICFP": "PCF",
        "PFCI": "FCI",
        "IIII": "",
        "FFFF": "CCCC",
        "CCCIIC": "III"}
    for pat, expected in tests.iteritems():
        testconsts(pat, expected)

def testnat(pat, expected):
    setdna(pat)
    if not isinstance(expected, int):
        try:
            n = nat()
            global dna
            print "Expected exception with pattern <%s>, got <%s> with dna <%s>" \
                % (pat, n, dna)
            raise AssertionError()
        except StopIteration: pass
    else:
        try:
            n = nat()
            if n != expected:
                print "Error in nat! expected <" + str(expected),
                print "> got: <" + str(n) + ">  on dna: <" + pat + ">"
                raise AssertionError()
        except Exception, e:
            print """Unexpected Exception in nat(); expected <%s> pat <%s>
    exception %s""" % (expected, pat, e)
            raise AssertionError()

def testnats():
    tests = {
        "": StopIteration,
        "P": 0,
        "IP": 0,
        "FP": 0,
        "CP": 1,
        "ICP": 2,
        "FCP": 2,
        "CCP": 3,
        "ICCP": 6,
        "IIIP": 0,
        "IFIP": 0,
        "IICCP": 12,
        "CICCP": 13,
    }
    for pat, expected in tests.iteritems():
        testnat(pat, expected)

def testpattern(pat, expected, exp_stack):
    global pqueue
    setdna(pat)
    p = pattern()
    for i in range(len(exp_stack)):
        if isinstance(exp_stack[i], str):
            exp_stack[i] = chararray(exp_stack[i])
    if p != chararray(expected) or pqueue != exp_stack:
        print """Error! expected <%s> got <%s>
exp_stack <%s> pqueue <%s>""" % (expected, p, exp_stack, pqueue)
        raise AssertionError()

def testpatterns():
    tests = [
        ("CIIC", "I", []),
        ("IIPIPICPIICICIIF", "(!)P", [2]),
        ("IIP"          #(
        +"IPICP"        #!nat() = 2
        +"IFCCCC"       #?consts() = "III"  //note consts consumes 3!
        +"IIC"          #)
        +"IFFFPIC"      #?consts() = "CFP"
        +"IIIFFFFFFF"   #add FFFFFFF to RNA
        +"IC"           #P
        +"IIF"          #return
        , "(!?)?P", ["CFP", "III", 2])

    ]
    for t in tests:
        testpattern(*t)

def generate_pattern(length, p_only_as_last = True):
    ret=""

    if p_only_as_last:
        letters = "ICF"
    else:
        letters = "ICFP"

    i = 0
    while i<length:
        ret += choice(letters)
        i += 1

    if p_only_as_last:
        ret += "P"

    return ret

def benchmark_quote():
    from time import time
    from sys import stdout
    print "Benchmarking quote():"
    print "1K elements: ", 
    stdout.flush()
    s = generate_pattern(1000, False)
    print "generated... ",
    stdout.flush()

    then = time()
    quote(s)
    print time()-then

    print "1M elements: ", 
    stdout.flush()
    s = generate_pattern(1000000, False)
    print "generated... ",
    stdout.flush()
    
    then = time()
    quote(s)
    print time()-then

    print "10M elements: ", 
    stdout.flush()
    s = generate_pattern(10000000, False)
    print "generated... ",
    stdout.flush()
    
    then = time()
    quote(s)
    print time()-then

if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1:
        if sys.argv[1] == "endo":
            f = file("endo.dna").read()
            execute(chararray(f))
        elif sys.argv[1] == "bench":
            f = file("endo.dna").read()[:100000]
            execute(chararray(f))
    else:
        try:
            testnats()
            testconstss()
            testpatterns()
            testtemplates()
            testquotes()
            testprotects()
            testasnats()
            testexecutes()
            print "\n*** All Tests Passed ***"
        except AssertionError:
            print "Tests Failed"
        except Exception, e:
            print "Exception rose too far!"
            print "dna <%s> rna <%s> pqueue <%s> tqueue <%s>" % (dna, rna, pqueue,
                tqueue)
            raise

        #benchmark_quote()
